import streamlit as st
from base import Agent, AgentLocal
from rag import query_doc
from rag import fix_json
from rag import ai_suggest_queries
from rag import clean_generated_text
from db import query_database
from db import extract_title
import json
import requests
import subprocess


def calculate_relevance_score(input_question, paper):
    prompt = '''你是一个非常厉害的科研助理。你会根据从知识库中搜集到的知识，帮我分析并回答问题。
    请根据生成知识库文章与问题之间的相关性得分，给出相关性得分。
    相关性得分范围为0-1，0表示完全不相关，1表示完全相关。
    相关性得分从研究的问题出发，考虑研究的问题与知识库文章的研究领域、研究问题、研究内容、研究方法、研究系统设定等的相似性。
    相关性得分越高，表示知识库文章与研究的问题越相关。
    '''
    agent1 = AgentLocal(prompt)
    score_txt = agent1.chat("用户的问题为："+input_question + "\n 请寻找论文与用户问题的相关性，并给出相关性得分：", 
                                context=f"<知识库检索> 论文内容：\n {paper} \n</知识库检索>")
    # 计算生成文本与问题之间的相关性得分
    # 这里可以根据具体需求使用不同的方法计算得分
    # 例如，可以使用余弦相似度计算生成文本与问题之间的相似性
    # 这里暂时返回一个示例得分
    return clean_generated_text(score_txt)


# 新增对话历史列表
history = []


def get_info(question, loops=1):   
    dbname = "arxiv"
    relevance_scores = None

    for i in range(loops):
        search_query = ai_suggest_queries(question, relevance_scores, json_loads=False, n=10)

        st.write("生成的检索问题: ", search_query)
        try:
            gen_txt, doc_ids = query_doc([search_query], dbname, n=100)
        except Exception as e:
            print(f"检索出错: {e}") 
            search_query = [question]
            gen_txt, doc_ids = query_doc(search_query, dbname, n=100)

        yield gen_txt, doc_ids

# pip install arxiv2bib
def fetch_bibtex(doc_id):
    # 这里添加获取 BibTeX 文件的逻辑
    # 假设我们从 arxiv 获取 BibTeX
    # url = f"https://arxiv.org/bibtex?id={doc_id}"
    try:
        # 使用 arxiv2bib 命令获取 BibTeX
        result = subprocess.run(['arxiv2bib', doc_id], capture_output=True, text=True, check=True)
        return result.stdout
    except subprocess.CalledProcessError as e:
        print(f"请求失败: {e}")
        return None

def get_sorted_bibtex(info):
    # 根据相关性得分排序
    #sorted_info = sorted(info, key=lambda x: x['score'], reverse=True)
    sorted_info = info
    bibtex_entries = []
    
    for item in sorted_info:
        doc_id = item['doc_id']
        bibtex = fetch_bibtex(doc_id)
        if bibtex:
            bibtex_entries.append(bibtex)
    
    return bibtex_entries




st.markdown(
    """
    <style>
    .reportview-container {
        background-color: #B0C4DE;  /* 灰蓝色 */
    }
    </style>
    """,
    unsafe_allow_html=True
)

st.title("科研助理")
input_question = st.text_area("请输入您的问题:", height=200)

if st.button("提交"):
    if input_question:
        info = []
        for papers, doc_ids in get_info(input_question, loops=1):

            gen_txt = ""
            
            # 新增一个集合用于存储已检索到的论文
            seen_papers = set()

            for i, paper in enumerate(papers):
                if paper not in seen_papers:
                    seen_papers.add(paper)  # 标记为已检索
                    title = extract_title(paper.replace("\n", ""))
                    query = query_database(title)
                    if len(query) == 0:
                        print(f"论文{title}没有检索到相关信息")

                    #gen_txt = agent.chat_local("用户的问题为："+input_question + "\n 请寻找论文与用户问题的相关性，并给出相关性得分：", context=f"<知识库检索> 论文内容：\n {paper} \n</知识库检索>")
                    reason = calculate_relevance_score(input_question, paper)
                    #score = calculate_relevance_score(gen_txt)  # 计算相关性得分
                    info.append({
                        'paper': paper,  # 保存论文信息
                        'analysis': reason,  # 保存分析内容
                        'query': query,  # 保存文档ID
                        'doc_id': query[0]
                    })
                    st.write(f"论文: {paper})")
                    st.write(f"信息: {query}")
                    st.write(reason)
                    st.write(f"人工智能解读：https://papers.cool/arxiv/{query[0]}")
                    st.write("-"*100)
            
            #info_txt = '\n'.join([item['analysis'] for item in info])  # 提取分析内容
            #gen_txt = agent.chat_local(input_question, context=f"<知识库检索> {info_txt} </知识库检索>")
            #agent.collect_message(gen_txt)

        # 去掉 info 中的重复元素
        unique_info = []
        seen_doc_ids = set()
        
        for item in info:
            doc_id = item['doc_id']
            if doc_id not in seen_doc_ids:
                seen_doc_ids.add(doc_id)
                unique_info.append(item)
        
        # 用去重后的列表替换原列表
        info = unique_info
        bibtex_list = get_sorted_bibtex(info)
        for bibtex in bibtex_list:
            st.write(bibtex)
        
    # 将 bibtex 列表输出到文件
    if bibtex_list:
        # 创建文件名
        file_name = "topic.bib"
        # 将所有 bibtex 条目写入文件
        with open(file_name, "w", encoding="utf-8") as f:
            for bibtex in bibtex_list:
                f.write(bibtex + "\n\n")
        
        st.success(f"已将参考文献列表保存到 {file_name} 文件中")
    else:
        st.warning("请先输入问题。")
